﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using SlimDX.Direct3D9;

namespace Gype.UI.D3D9
{
    internal class ShaderCompiler<T>
    {
        public ShaderCompiler(Device dev, string code, string entryPoint, string profile)
        {
            mDevice = dev;
            Code = ShaderBytecode.Compile(code, entryPoint, profile, ShaderFlags.SkipOptimization | ShaderFlags.Debug);
            Shader = (T)Activator.CreateInstance(typeof(T), dev, Code);
            ConstantTable = new Dictionary<string, ConstantDescription>();

            reflectShader();
        }

        public ShaderCompiler(Device dev, byte[] code, string entryPoint, string profile)
            : this(dev, Encoding.ASCII.GetString(code), entryPoint, profile)
        {
        }

        public void SetValue(ConstantDescription desc, Texture tex)
        {
            mDevice.SetTexture(desc.RegisterIndex, tex);
        }

        public void SetValue(ConstantDescription desc, int value)
        {
            if (Shader is VertexShader)
            {
                VertexShader oldShader = null;
                if (mDevice.VertexShader != (VertexShader)(object)Shader)
                {
                    oldShader = mDevice.VertexShader;
                    mDevice.VertexShader = (VertexShader)(object)Shader;
                }

                mDevice.SetVertexShaderConstant(desc.RegisterIndex, new int[] { value });

                if (oldShader != null)
                    mDevice.VertexShader = oldShader;
            }
            else if (Shader is PixelShader)
            {
                mDevice.SetPixelShaderConstant(desc.RegisterIndex, new int[] { value });
            }
        }

        public void SetValue(ConstantDescription desc, LinearMatrix matrix)
        {
            if (Shader is VertexShader)
            {
                VertexShader oldShader = null;
                if (mDevice.VertexShader != (VertexShader)(object)Shader)
                {
                    oldShader = mDevice.VertexShader;
                    mDevice.VertexShader = (VertexShader)(object)Shader;
                }

                mDevice.SetVertexShaderConstant(desc.RegisterIndex, new float[]
                    {
                        matrix.M11, matrix.M12, matrix.M13, 0
                    }
                    );

                mDevice.SetVertexShaderConstant(desc.RegisterIndex + 1, new float[]
                {
                    matrix.M21, matrix.M22, matrix.M23, 0
                }
                );

                mDevice.SetVertexShaderConstant(desc.RegisterIndex + 2, new float[]
                {
                    matrix.M31, matrix.M32, matrix.M33, 0
                }
                );

                if (oldShader != null)
                    mDevice.VertexShader = oldShader;
            }
            else if (Shader is PixelShader)
            {
                mDevice.SetVertexShaderConstant(desc.RegisterIndex, new float[]
                    {
                        matrix.M11, matrix.M12, matrix.M13,
                    }
                    );

                mDevice.SetVertexShaderConstant(desc.RegisterIndex + 1, new float[]
                {
                    matrix.M21, matrix.M22, matrix.M23,
                }
                );

                mDevice.SetVertexShaderConstant(desc.RegisterIndex + 2, new float[]
                {
                    matrix.M31, matrix.M32, matrix.M33
                }
                );
            }
        }

        public void SetValue(ConstantDescription desc, SlimDX.Matrix value)
        {
            if (Shader is VertexShader)
            {
                VertexShader oldShader = null;
                if (mDevice.VertexShader != (VertexShader)(object)Shader)
                {
                    oldShader = mDevice.VertexShader;
                    mDevice.VertexShader = (VertexShader)(object)Shader;
                }

                mDevice.SetVertexShaderConstant(desc.RegisterIndex, value);

                if (oldShader != null)
                    mDevice.VertexShader = oldShader;
            }
            else if (Shader is PixelShader)
            {
                PixelShader oldShader = null;
                if (mDevice.PixelShader != (PixelShader)(object)Shader)
                {
                    oldShader = mDevice.PixelShader;
                    mDevice.PixelShader = (PixelShader)(object)Shader;
                }

                mDevice.SetPixelShaderConstant(desc.RegisterIndex, value);

                if (oldShader != null)
                    mDevice.PixelShader = oldShader;
            }
        }

        public void SetValue(ConstantDescription desc, float value)
        {
            if (Shader is VertexShader)
            {
                VertexShader oldShader = null;
                if (mDevice.VertexShader != (VertexShader)(object)Shader)
                {
                    oldShader = mDevice.VertexShader;
                    mDevice.VertexShader = (VertexShader)(object)Shader;
                }

                mDevice.SetVertexShaderConstant(desc.RegisterIndex, new SlimDX.Vector4[] { new SlimDX.Vector4(value, 0, 0, 0) });

                if (oldShader != null)
                    mDevice.VertexShader = oldShader;
            }
            else if (Shader is PixelShader)
            {
                PixelShader oldShader = null;
                if (mDevice.PixelShader != (PixelShader)(object)Shader)
                {
                    oldShader = mDevice.PixelShader;
                    mDevice.PixelShader = (PixelShader)(object)Shader;
                }

                mDevice.SetPixelShaderConstant(desc.RegisterIndex, new SlimDX.Vector4[] { new SlimDX.Vector4(value, 0, 0, 0) });

                if (oldShader != null)
                    mDevice.PixelShader = oldShader;
            }
        }

        public void SetValue(ConstantDescription desc, SlimDX.Vector2 value)
        {
            if (Shader is VertexShader)
            {
                VertexShader oldShader = null;
                if (mDevice.VertexShader != (VertexShader)(object)Shader)
                {
                    oldShader = mDevice.VertexShader;
                    mDevice.VertexShader = (VertexShader)(object)Shader;
                }

                mDevice.SetVertexShaderConstant(desc.RegisterIndex, new SlimDX.Vector4[] { new SlimDX.Vector4(value.X, value.Y, 0, 0) });

                if (oldShader != null)
                    mDevice.VertexShader = oldShader;
            }
            else if (Shader is PixelShader)
            {
                PixelShader oldShader = null;
                if (mDevice.PixelShader != (PixelShader)(object)Shader)
                {
                    oldShader = mDevice.PixelShader;
                    mDevice.PixelShader = (PixelShader)(object)Shader;
                }

                mDevice.SetPixelShaderConstant(desc.RegisterIndex, new SlimDX.Vector4[] { new SlimDX.Vector4(value.X, value.Y, 0, 0) });

                if (oldShader != null)
                    mDevice.PixelShader = oldShader;
            }
        }

        public void SetValue(ConstantDescription desc, SlimDX.Vector3 value)
        {
            if (Shader is VertexShader)
            {
                mDevice.VertexShader = (VertexShader)(object)Shader;
                mDevice.SetVertexShaderConstant(desc.RegisterIndex, new SlimDX.Vector4[] { new SlimDX.Vector4(value.X, value.Y, value.Z, 0) });
                mDevice.VertexShader = null;
            }
            else if (Shader is PixelShader)
            {
                PixelShader oldShader = null;
                if (mDevice.PixelShader != (PixelShader)(object)Shader)
                {
                    oldShader = mDevice.PixelShader;
                    mDevice.PixelShader = (PixelShader)(object)Shader;
                }

                mDevice.SetPixelShaderConstant(desc.RegisterIndex, new SlimDX.Vector4[] { new SlimDX.Vector4(value.X, value.Y, value.Z, 0) });

                if (oldShader != null)
                    mDevice.PixelShader = oldShader;
            }
        }

        public void SetValue(ConstantDescription desc, SlimDX.Vector4 value)
        {
            if (Shader is VertexShader)
            {
                mDevice.VertexShader = (VertexShader)(object)Shader;
                mDevice.SetVertexShaderConstant(desc.RegisterIndex, new SlimDX.Vector4[] { value });
                mDevice.VertexShader = null;
            }
            else if (Shader is PixelShader)
            {
                PixelShader oldShader = null;
                if (mDevice.PixelShader != (PixelShader)(object)Shader)
                {
                    oldShader = mDevice.PixelShader;
                    mDevice.PixelShader = (PixelShader)(object)Shader;
                }

                mDevice.SetPixelShaderConstant(desc.RegisterIndex, new SlimDX.Vector4[] { value });

                if (oldShader != null)
                    mDevice.PixelShader = oldShader;
            }
        }

        private void reflectShader()
        {
            var numConstants = Code.ConstantTable.Description.Constants;
            for (int i = 0; i < numConstants; ++i)
            {
                var handle = Code.ConstantTable.GetConstant(null, i);
                var desc = Code.ConstantTable.GetConstantDescription(handle);
                ConstantTable.Add(desc.Name, desc);
            }
        }

        public ShaderBytecode Code { get; private set; }
        public T Shader { get; private set; }
        public Dictionary<string, ConstantDescription> ConstantTable { get; private set; }

        private Device mDevice;
    }
}
